from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import sys
import numpy as np
import onnx
from onnx import helper
from helper import make_model, create_onnx_session
from onnx import AttributeProto, TensorProto, GraphProto

def onnx_pad_node_example(op_type="", node_index=1):
    # The protobuf definition can be found here:
    # https://github.com/onnx/onnx/blob/master/onnx/onnx.proto

    # Create one input (ValueInfoProto)
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 2])
    # Create second input (ValueInfoProto)
    pads = helper.make_tensor_value_info('pads', TensorProto.INT64, [4])
    # Create third input (ValueInfoProto)
    values = helper.make_tensor_value_info('values', TensorProto.FLOAT, [1])
    # Create one output (ValueInfoProto)
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 4])

    # Create a node (NodeProto)
    node_def = helper.make_node(
        op_type=op_type,
        name='%s_%d' % (op_type, node_index),
        inputs=['x', 'pads', 'values'], # inputs
        outputs=['y'], # outputs
        mode='constant', # Attributes
    )

    # Create the graph (GraphProto)
    graph_def = helper.make_graph(
        [node_def],
        "test-model",
        [x, pads, values],
        [y],
    )

    ## make onnx model file
    model_name = '%s.onnx'% op_type.lower()
    make_model(graph_def, model_name)
    ## Start onnx runtime
    sess = create_onnx_session(model_name)

    ## Generate input data
    x = np.random.randn(1, 2).astype(np.float32)
    pads = np.array([0, 1, 1, 2]).astype(np.int64)  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
    values = np.array([1.2]).astype(np.float32)
    print("input:", x)

    inputs = {"x":x, "pads":pads, "values":values}
    outputs = ['y']
    result = sess.run(outputs, inputs)
    print(result[0].shape)

def onnx_scatterelements_node_example(op_type="", node_index=1):
    bs = 10
    size_1 = 128
    size_2 = 64

    # The protobuf definition can be found here:
    # https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
    # Create one input (ValueInfoProto)
    data = helper.make_tensor_value_info('data', TensorProto.FLOAT, [bs, size_1])
    # Create second input (ValueInfoProto)
    indices = helper.make_tensor_value_info('indices', TensorProto.INT64, [bs, size_2])
    # Create third input (ValueInfoProto)
    updates = helper.make_tensor_value_info('updates', TensorProto.FLOAT, [bs, size_2])
    axis = 1
    # Create one output (ValueInfoProto)
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [bs, size_1])

    # Create a node (NodeProto)
    node_def = onnx.helper.make_node(
        op_type=op_type,
        name='%s_%d' % (op_type, node_index),
        inputs=['data', 'indices', 'updates'],
        outputs=['y'],
        axis=axis,
    )

    # Create the graph (GraphProto)
    graph_def = helper.make_graph(
        [node_def],
        "test-model",
        [data, indices, updates],
        [y],
    )

    ## make onnx model file
    model_name = '%s.onnx'% op_type.lower()
    make_model(graph_def, model_name)
    ## Start onnx runtime
    sess = create_onnx_session(model_name)

    ## Generate input data
    x = np.zeros((bs, size_1), dtype=np.float32)
    index = np.zeros((bs, size_2), dtype=np.int64)
    for i in range(bs):
        for j in range(10):
            index[i][j] = np.random.randint(low=0, high=size_1)
    values = np.zeros((bs, size_2), dtype=np.float32)
    for i in range(bs):
        for j in range(10):
            values[i][j] = np.random.rand()

    inputs = {"data":x, "indices":index, "updates":values}
    output = ['y']
    result = sess.run(output, inputs)
    print(result[0].shape)

def onnx_three_nodes_example(op_type="", node_index=1):
    bs = 1000
    size = 21128

    # The protobuf definition can be found here:
    # https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
    # Create one input (ValueInfoProto)
    x = helper.make_tensor_value_info('x', TensorProto.INT32, [bs, size])

    # Create NonZero node (NodeProto)
    node_index = node_index + 1
    op_type = 'NonZero'
    nonzero_name = '%s_%d' % (op_type, node_index)
    nonzero_node = onnx.helper.make_node(
        op_type=op_type,
        name=nonzero_name,
        inputs=['x'],
        outputs=[nonzero_name],
    )

    # Create Cast node (NodeProto)
    node_index = node_index + 1
    op_type = 'Cast'
    cast_name = '%s_%d' % (op_type, node_index)
    cast_node = onnx.helper.make_node(
        op_type=op_type,
        name=cast_name,
        inputs=[nonzero_name],
        outputs=[cast_name],
        to=getattr(TensorProto, 'FLOAT'),
    )

    # Create one output (ValueInfoProto)
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10000, 2])

    # Create Transpose node (NodeProto)
    node_index = node_index + 1
    op_type = 'Transpose'
    transpose_name = '%s_%d' % (op_type, node_index)
    transpose_node = onnx.helper.make_node(
        op_type=op_type,
        name=transpose_name,
        inputs=[cast_name],
        outputs=['y'],
        perm=[1,0],
    )

    # Create the graph (GraphProto)
    graph_def = helper.make_graph(
        [nonzero_node, cast_node, transpose_node],
        "test-model",
        [x],
        [y],
    )

    ## make onnx model file
    model_name = 'nonzero_cast_transpose.onnx'
    make_model(graph_def, model_name)
    ## Start onnx runtime
    sess = create_onnx_session(model_name)

    ## Generate input data
    data = np.zeros((bs, size), dtype=np.int32)
    for i in range(bs):
        for j in range(10):
            col = np.random.randint(low=0, high=size)
            value = np.random.randint(low=0, high=size)
            data[i][col] = value
   
    inputs = {"x":data}
    output = ['y']
    result = sess.run(output, inputs)
    print(result[0].shape)

def main():
    global_index = 0
    ## Build 'Pad' node
    global_index = global_index + 1
    onnx_pad_node_example(op_type = "Pad", node_index=global_index)

    ## Build 'ScatterElements' node
    global_index = global_index + 1
    onnx_scatterelements_node_example(op_type = "ScatterElements", node_index=global_index)

    ## Build 'Nonzero/Cast/Transpose' node. Very useful example!
    global_index = global_index + 1
    onnx_three_nodes_example(op_type = "", node_index=global_index)

## [W:onnxruntime:, graph.cc:73 MergeShapeInfo] Error merging shape info for output. 
## 'y' source:{10,128} target:{1,4}. Falling back to lenient merge.
## the output shape is not match the result shape 

if __name__ == "__main__":
    main()
